package com.tgy.tree;

import java.util.Comparator;

/**
 * @Author: tgy
 * @Date: 2020-09-18 09:30
 */
public class RBTree<E> extends BinarySearchTree<E>{

    public RBTree(Comparator<E> comparator) {
        super(comparator);
    }

    public RBTree() {

        this(null);
    }

    private static enum RotateEnum{
        /**
         * 向左旋转
         */
        LEFT,
        /**
         * 向右旋转
         */
        RIGHT

    }

    private static enum Color{

        /**
         * 黑色节点
         */
        BLACK,
        /**
         * 红色节点
         */
        RED

    }

    private static class RBNode<E> extends Node<E> {

        /**
         * 默认节点是红色
         */
        private Color color = Color.RED;

        public RBNode(E element, Node<E> parent) {
            super(element, parent);
        }

        /**
         * 判断当前节点是否是黑色节点
         * @return
         */
        public boolean isBlack() {

            return color == Color.BLACK;
        }

        /**
         * 判断当前节点是否是红色节点
         * @return
         */
        public boolean isRed() {

            return color == Color.RED;
        }

        /**
         * 把当前节点变成黑色
         * @return
         */
        public RBNode<E> black(){

            color = Color.BLACK;
            return this;
        }

        /**
         * 把当前节点变成红色
         * @return
         */
        public RBNode<E> red() {

            color = Color.RED;

            return this;
        }

        public RBNode<E> color(Color color) {

            this.color = color;
            return this;
        }

        /**
         * 获取父节点
         * @return
         */
        public RBNode<E> parent(){

            return (RBNode<E>) this.parent;
        }

        /**
         * 获取左节点
         * @return
         */
        public RBNode<E> leftNode() {

            return (RBNode<E>) this.left;
        }

        /**
         * 获取右节点
         * @return
         */
        public RBNode<E> rightNode() {

            return (RBNode<E>)this.right;
        }



        /**
         * 获取当前节点的兄弟节点
         * @return
         */
        public RBNode<E> sibling() {

            if (parent == null) {

                return null;
            }

            if (isLeftNode()) {

                return parent().rightNode();
            }else if (isRightNode()) {

                return parent().leftNode();

            }
            return null;
        }

        /**
         * 获取父节点的兄弟节点
         * @return
         */
        public RBNode<E> uncle() {

            if (grand() == null) {

                return null;
            }

            if (parent().isLeftNode()) {

                return grand().rightNode();
            }else{

                return grand().leftNode();
            }
        }



        /**
         * 获取祖父节点(父节点的父节点)
         * @return
         */
        public RBNode<E> grand() {

            return parent() == null? null: parent().parent();
        }

        @Override
        public String toString() {

            return element +":(" + color.name() + ")";
        }

        public static boolean isBlack(RBNode node) {

            if (node == null) {

                return true;
            }

            return node.isBlack();
        }

    }

    @Override
    protected Node<E> createNode(E element, Node<E> parentNode) {

        return new RBNode<>(element,parentNode);
    }

    @Override
    protected void addAdjustTree(Node<E> node) {

        RBNode<E> rbNode = (RBNode<E>) node;

        /**
         * 根节点
         */
        if (rbNode.parent == null) {

            rbNode.black();
            return;
        }

        /**
         * 如果添加节点的父节点是黑色，则直接添加，什么都不做
         */
        if (rbNode.parent().isBlack()) {

            return;
        }


        if (RBNode.isBlack(rbNode.uncle())) {

            RBNode<E> grand = rbNode.grand();
            RBNode<E> parent = rbNode.parent();
            // 如果uncle 为黑色,旋转
            if (parent.isLeftNode()) {

//                L
                if (rbNode.isLeftNode()) {

//                    L
                      parent.black();
                }else {
//                    R
                    rotate(parent,RotateEnum.LEFT);
//                    rotate(grand,RotateEnum.RIGHT);
                    rbNode.black();
                }
                rotate(grand,RotateEnum.RIGHT);

            }else {

//                R
                if (rbNode.isLeftNode()) {
//                    L
                    rotate(parent,RotateEnum.RIGHT);
                    rbNode.black();
                }else {
//                    R
//                    rotate(grand,RotateEnum.LEFT);
                    parent.black();
                }
                rotate(grand,RotateEnum.LEFT);
            }

            grand.red();
        }else {

            // 拆分B树节点
            rbNode.uncle().black();
            rbNode.grand().red();
            rbNode.parent().black();
            addAdjustTree(rbNode.grand());
        }
    }

    @Override
    protected void removeAdjustTree(Node<E> node) {

        RBNode<E> ebNode = (RBNode<E>) node;

        if (ebNode.isRed() || node.parent == null) {

            ebNode.black();
            return;
        }

        RBNode<E> parentNode = ebNode.parent();

        /**
         *      由于在BinarySearchTree 中如果是左树，则把父节点的left置为null 判断是左树还是右树
         *      1. parentNode.leftNode() == null
         *          从父类进入本方法
         *      2. node.isLeftNode()
         *          从下面321行的模拟父节点删除进入本方法
         *
         */
        boolean isLeft = (parentNode.leftNode() == null) || node.isLeftNode();

        if (isLeft) {
            // 判断兄弟节点是否是black
            RBNode<E> sibling = parentNode.rightNode();
            if (sibling.isRed()) {

                rotate(node.parent,RotateEnum.LEFT);
                sibling.black();
                parentNode.red();
                sibling = parentNode.rightNode();
            }

            if (RBNode.isBlack(sibling.leftNode())  && RBNode.isBlack(sibling.rightNode())) {

                // 兄弟节点染成红色
                sibling.red();
                // 没有红色的子节点
                if (parentNode.isRed()) {
//                      父节点是红色,直接把父节点染成黑色
                    parentNode.black();
                }else {
//                     相当于删除父节点,重新走一遍调整逻辑
                    removeAdjustTree(parentNode);
                }
            }else {

                // 如果兄弟节点的右节点是黑色，则左节点一定是红色，则进行右旋转
                if (RBNode.isBlack(sibling.rightNode())) {

                    rotate(sibling,RotateEnum.RIGHT);
                    sibling.red();
                    sibling =parentNode.rightNode();
                }

                sibling.color(parentNode.color);
                parentNode.black();
                if (sibling.rightNode() != null) {

                    sibling.rightNode().black();
                }
                rotate(ebNode.parent(),RotateEnum.LEFT);
            }

        }else {
        // 判断兄弟节点是否是black
            RBNode<E> sibling = parentNode.leftNode();
            if (sibling.isRed()) {

                rotate(node.parent,RotateEnum.RIGHT);
                sibling.black();
                parentNode.red();
                sibling = parentNode.leftNode();

            }

            if (RBNode.isBlack(sibling.leftNode())  && RBNode.isBlack(sibling.rightNode())) {

                // 兄弟节点染成红色
                sibling.red();
                  // 没有红色的子节点
              if (parentNode.isRed()) {
//                      父节点是红色,直接把父节点染成黑色
                  parentNode.black();
              }else {
//                     相当于删除父节点,重新走一遍调整逻辑
                  removeAdjustTree(parentNode);
              }
            }else {

                // 如果兄弟节点的左字节是黑色，则右节点一定是红色，则进行左旋转
                if (RBNode.isBlack(sibling.leftNode())) {

                    rotate(sibling,RotateEnum.LEFT);
                    sibling.red();
                    sibling = parentNode.leftNode();
                }

                sibling.color(parentNode.color);
                parentNode.black();
                if (sibling.leftNode() != null) {

                    sibling.leftNode().black();
                }
                rotate(ebNode.parent(),RotateEnum.RIGHT);
            }
        }

    }

    private void rotate(Node<E> gNode, RotateEnum rotateEnum){

        Node<E> pNode;


        if (rotateEnum == RotateEnum.LEFT) {

            pNode = gNode.right;
        }else {

            pNode = gNode.left;
        }

        if (gNode.parent == null) {

            root = pNode;
        }else {

            if (gNode.isLeftNode()){

                gNode.parent.left = pNode;
            }else {

                gNode.parent.right = pNode;
            }
        }

        pNode.parent = gNode.parent;
        gNode.parent = pNode;

        if (rotateEnum == RotateEnum.LEFT) {

            if (pNode.left != null) {

                pNode.left.parent = gNode;
            }
            gNode.right = pNode.left;
            pNode.left = gNode;

        }else {

            if (pNode.right != null) {

                pNode.right.parent = gNode;
            }
            gNode.left = pNode.right;
            pNode.right = gNode;
        }
    }
}
